/*
 * Copyright (c) 2021 Futurewei Technologies, Inc.
 *
 * clang2mpl is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan
 * PSL v2. You may obtain a copy of Mulan PSL v2 at:
 *
 *     http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
 * KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
 * NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. See the
 * Mulan PSL v2 for more details.
 */
#pragma once

#include <map>
#include <stack>

#include "Clang2MapleResult.h"
#include "clang/AST/DeclVisitor.h"
#include "clang/AST/StmtVisitor.h"
#include "clang/AST/Type.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Frontend/FrontendAction.h"
#include "clang/Tooling/Tooling.h"
#include "llvm/Support/raw_ostream.h"

// For Maple IR
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunknown-warning-option"
#pragma GCC diagnostic ignored "-Wcast-qual"
#pragma GCC diagnostic ignored "-Wignored-qualifiers"
#pragma GCC diagnostic ignored "-Wreturn-type"
#pragma GCC diagnostic ignored "-Woverloaded-virtual"
#pragma GCC diagnostic ignored "-Wunused-variable"
#pragma GCC diagnostic ignored "-Wsign-compare"
#pragma GCC diagnostic ignored "-Wpedantic"
#pragma GCC diagnostic ignored "-Wtype-limits"
#pragma GCC diagnostic ignored "-Wcovered-switch-default"
#pragma GCC diagnostic ignored "-Wsuggest-override"
#include "mir_builder.h"
#include "mir_nodes.h"
#include "mir_type.h"
#pragma GCC diagnostic pop

using namespace maple;

struct LabelInfo {
  LabelIdx Label;
  bool Used = false;

  LabelInfo(LabelIdx LI) : Label(LI) {}
};

class Clang2MapleVisitor
    : public clang::ConstStmtVisitor<Clang2MapleVisitor, Result>,
      public clang::ConstDeclVisitor<Clang2MapleVisitor, MIRSymbol *> {
private:
  clang::ASTContext *Context = nullptr;
  MIRBuilder *Builder = nullptr;
  MIRModule *Module = nullptr;
  uint32_t UnnamedCount = 0;
  int FileIdx = 2;
  // This flag is set to true when visiting an initializer for a global or
  // static local variable. For an initializer, the type of a constant must be
  // its true type, but for other expressions, it should be a register-sized
  // type. This flag enables that behavior.
  bool Initializer = false;
  std::map<const clang::Type *, TyIdx> TypeMap;
  std::map<const clang::RecordDecl *, MIRPtrType *> UnresolvedTypes;
  std::map<unsigned, int> FileMap;
  std::map<const clang::Decl *, MIRSymbol *> DeclMap;
  std::stack<LabelInfo> EndLabels;
  std::stack<LabelInfo> ExitLabels;
  std::stack<LabelInfo> DefaultLabels;
  std::stack<CaseVector *> SwitchCases;
  // FieldMap[TopLevelRecord][BaseFieldID][FieldDecl]
  std::map<const clang::RecordDecl *,
           std::map<FieldID, std::map<const clang::FieldDecl *, FieldID>>>
      FieldMap;
  PrimType PointerPrimTy;
  MIRType *IntPointerTy;

public:
  using clang::ConstStmtVisitor<Clang2MapleVisitor, Result>::Visit;
  using clang::ConstDeclVisitor<Clang2MapleVisitor, MIRSymbol *>::Visit;

  Clang2MapleVisitor() {}

  MIRSymbol *VisitTranslationUnitDecl(const clang::TranslationUnitDecl *TU);

  MIRSymbol *VisitTypeDecl(const clang::TypeDecl *Ty);
  MIRSymbol *VisitRecordDecl(const clang::RecordDecl *Record);
  MIRSymbol *VisitTypedefDecl(const clang::TypedefDecl *Typedef);
  MIRSymbol *VisitVarDecl(const clang::VarDecl *Var);
  MIRSymbol *VisitParmVarDecl(const clang::ParmVarDecl *Var);
  MIRSymbol *VisitFunctionDecl(const clang::FunctionDecl *Func);
  MIRSymbol *VisitEnumDecl(const clang::EnumDecl *Enum);
  MIRSymbol *VisitEnumConstantDecl(const clang::EnumConstantDecl *EnumConst);
  MIRSymbol *VisitFileScopeAsmDecl(const clang::FileScopeAsmDecl *AsmDecl);

  // This method is called first by the visitor for all statements. It checks if
  // an expression can be simplified into a constant, and if not, then it
  // traverses the tree.
  Result Visit(const clang::Stmt *S);

  Result VisitStmt(const clang::Stmt *S);
  Result VisitGCCAsmStmt(const clang::GCCAsmStmt *GAS);
  Result VisitCompoundStmt(const clang::CompoundStmt *CS);
  Result VisitDeclStmt(const clang::DeclStmt *DS);
  Result VisitDoStmt(const clang::DoStmt *Do);
  Result VisitForStmt(const clang::ForStmt *For);
  Result VisitIfStmt(const clang::IfStmt *If);
  Result VisitNullStmt(const clang::NullStmt *N);
  Result VisitReturnStmt(const clang::ReturnStmt *Return);
  Result VisitWhileStmt(const clang::WhileStmt *While);
  Result VisitBreakStmt(const clang::BreakStmt *BS);
  Result VisitContinueStmt(const clang::ContinueStmt *CS);
  Result VisitLabelStmt(const clang::LabelStmt *LS);
  Result VisitGotoStmt(const clang::GotoStmt *Goto);
  Result VisitIndirectGotoStmt(const clang::IndirectGotoStmt *Goto);
  Result VisitCaseStmt(const clang::CaseStmt *CS);
  Result VisitDefaultStmt(const clang::DefaultStmt *DS);
  Result VisitSwitchStmt(const clang::SwitchStmt *Switch);

  Result VisitAddrLabelExpr(const clang::AddrLabelExpr *ALE);
  Result VisitArraySubscriptExpr(const clang::ArraySubscriptExpr *ASE);
  Result VisitBinaryOperator(const clang::BinaryOperator *BO);
  Result
  VisitBinaryConditionalOperator(const clang::BinaryConditionalOperator *BCO);
  Result VisitCallExpr(const clang::CallExpr *CE);
  Result VisitCastExpr(const clang::CastExpr *CE);
  Result VisitCharacterLiteral(const clang::CharacterLiteral *CL);
  Result VisitChooseExpr(const clang::ChooseExpr *CE);
  Result VisitCompoundLiteralExpr(const clang::CompoundLiteralExpr *CLE);
  Result VisitConditionalOperator(const clang::ConditionalOperator *CO);
  Result VisitConstantExpr(const clang::ConstantExpr *CE);
  Result VisitDeclRefExpr(const clang::DeclRefExpr *DRE);
  Result VisitFloatingLiteral(const clang::FloatingLiteral *Lit);
  Result VisitImplicitValueInitExpr(const clang::ImplicitValueInitExpr *IV);
  Result VisitInitListExpr(const clang::InitListExpr *ILE);
  Result VisitIntegerLiteral(const clang::IntegerLiteral *Lit);
  Result VisitMemberExpr(const clang::MemberExpr *ME);
  Result VisitOffsetOfExpr(const clang::OffsetOfExpr *OOE);
  Result VisitOpaqueValueExpr(const clang::OpaqueValueExpr *OVE);
  Result VisitParenExpr(const clang::ParenExpr *PE);
  Result VisitPredefinedExpr(const clang::PredefinedExpr *PE);
  Result VisitShuffleVectorExpr(const clang::ShuffleVectorExpr *SVE);
  Result VisitStmtExpr(const clang::StmtExpr *SE);
  Result VisitStringLiteral(const clang::StringLiteral *Lit);
  Result VisitTypeTraitExpr(const clang::TypeTraitExpr *TTE);
  Result
  VisitUnaryExprOrTypeTraitExpr(const clang::UnaryExprOrTypeTraitExpr *UETT);
  Result VisitUnaryOperator(const clang::UnaryOperator *UO);
  Result VisitVAArgExpr(const clang::VAArgExpr *VAArg);

private:
  void setupBuiltinTypes(void);
  TyIdx type2MplIdx(clang::QualType QT, bool needComplete = true);
  MIRType *type2Mpl(clang::QualType QT);
  int fileID2Mpl(clang::FileID FID);
  MIRSymbol *decl2Mpl(const clang::Decl *D);
  SrcPosition sourceLocation2Mpl(clang::SourceLocation Loc);
  int typeRank(PrimType PTy);
  PrimType typeOfOperation(PrimType PTy1, PrimType PTy2);
  void addStmtToBlock(BlockNode *Block, StmtNode *SNode);
  void addToBlock(BlockNode *Block, Result &Res);
  BaseNode *exprToCond(BaseNode *Cond);
  BaseNode *exprToNotCond(BaseNode *Cond);
  FieldID recordFieldIDs(const clang::RecordDecl *Record,
                         const clang::RecordDecl *Ancestor, FieldID Base);
  BaseNode *getNodeAsLVal(Result &Res);
  BaseNode *getNodeAsRVal(Result &Res);
  StmtNode *getNodeAsStmt(Result &Res);
  BaseNode *getNodeAsAddrOf(Result &Res);
  void appendResultBefore(Result &To, Result R);
  bool evaluateNodeAsInt(int64_t &Result, const BaseNode *N);
  MIRConst *evaluateExprAsConst(const clang::Expr *E, MIRType *Ty);
  MIRConst *createZero(MIRType *Ty);
  BaseNode *createVectorZero(MIRType *Ty);
  Result buildExprToComputeSizeFromVLA(const clang::Expr *E,
                                       const clang::QualType &QT);
  BaseNode *constToNode(MIRConst *Const);
  unsigned assignStruct(Result &Res, MIRSymbol *BaseSym, BaseNode *BaseAddr,
                        MIRStructType *StructTy, BaseNode *Src,
                        FieldID Base = 0);
  unsigned assignArray(Result &Res, BaseNode *BaseAddr, MIRArrayType *StructTy,
                       BaseNode *Src);
  unsigned assignVector(Result &Res, MIRSymbol *VecSym, MIRType *Ty,
                        BaseNode *Src);
  StmtNode *vectorSetLane(MIRType *Ty, Result &VecRes, BaseNode *Index,
                          BaseNode *Src);
  PrimType getVectorElementPrimType(PrimType VectorPrimType);
  MIRConst *getInitializer(MIRType *Ty, const clang::Expr *InitExpr);
  MIRConst *expandConstants(MIRType *Ty, MIRConst *Const);
  MIRConst *completeArrayInitializer(MIRAggConst *Agg, MIRArrayType *ArrTy);
  MIRType *isHomogenousAggregate(MIRType *Ty);
  bool passInGeneralRegisters(MIRType *Ty);
  bool passInFloatRegisters(MIRType *Ty);
  BaseNode *getBase(MIRSymbol *Sym, BaseNode *Base, FieldID Field);
  AddrofNode *computeArrayOffset(size_t &Offset, ArrayNode *Array);
  BaseNode *getAddrOfNode(BaseNode *N);
  // Read the specified field, handling either symbol or node base
  BaseNode *readField(Result BaseRes, FieldID Field, MIRType *Ty);
  // Write the specified field, handling either symbol or node base
  StmtNode *writeField(Result BaseRes, FieldID Field, BaseNode *Src);
  bool shouldEvaluateArgs(const clang::FunctionDecl *FD);
  bool isNoCvtNeeded(PrimType ToPTy, PrimType FromPTy);
  bool isPointerIntConvert(PrimType ToPTy, PrimType FromPTy);
  BaseNode *createNumericCast(MIRType *ToTy, BaseNode *N);
};

class Clang2MapleConsumer : public clang::ASTConsumer {
public:
  Clang2MapleConsumer() : Visitor() {}

  virtual void HandleTranslationUnit(clang::ASTContext &Context) override {
    Visitor.Visit(Context.getTranslationUnitDecl());
  }

private:
  Clang2MapleVisitor Visitor;
};

class Clang2MapleAction : public clang::ASTFrontendAction {
public:
  Clang2MapleAction() {}

  virtual std::unique_ptr<clang::ASTConsumer>
  CreateASTConsumer(clang::CompilerInstance &Compiler,
                    llvm::StringRef InFile) override {
    return std::unique_ptr<clang::ASTConsumer>(new Clang2MapleConsumer());
  }
};

std::unique_ptr<clang::tooling::FrontendActionFactory>
clang2MapleFrontendActionFactory();
